"""
Author: Morphlng
Date: 2024-03-18 16:09:30
LastEditTime: 2025-02-15 16:36:53
LastEditors: Morphlng
Description: Scenario configuration helper
FilePath: /DrivingGym/src/driving_gym/environment/scenario/config/scenario_config.py
"""

import copy


class RoutingTask:
    def __init__(self, start=None, end=None, random=False):
        self._start = start or (0, 0, 0)
        self._end = end or (0, 0, 0)
        self._route = []
        self._init_speed = 0
        self._random = random
        self._min_distance = -1
        self._max_distance = float("inf")

    def starting(self, start):
        """Set the starting position of the task"""
        self._start = start
        return self

    def ending(self, end):
        """Set the ending position of the task"""
        self._end = end
        return self

    def routing(self, route):
        """Set the routing of the task. If start and end are set, the route will be appended between them"""
        if self._start != (0, 0, 0):
            self._route.append(self._start)
        self._route.extend(route)
        if self._end != (0, 0, 0):
            self._route.append(self._end)
        return self

    def initial_speed(self, speed):
        self._init_speed = speed
        return self

    def randomize(self, random=True, min_distance=-1, max_distance=float("inf")):
        self._random = random
        self._min_distance = min_distance
        self._max_distance = max_distance
        return self

    def to_dict(self):
        return {
            "type": "routing",
            "start": self._start,
            "end": self._end,
            "route": self._route,
            "random": self._random,
            "init_speed": self._init_speed,
            "min_distance": self._min_distance,
            "max_distance": self._max_distance,
        }


class ScenarioConfig:
    def __init__(self, scenario_type: str = "base"):
        self._type = scenario_type
        self._max_steps = 500
        self._map = ""
        self._actor = {}
        self._task = {}

    def environment(self, map_name: str, max_steps: int = 500):
        self._map = map_name
        self._max_steps = max_steps
        return self

    def add_task(self, actor_name: str, task: dict):
        """Add a task to the scenario

        Args:
            actor_name (str): The actor to perform the task
            task (dict): The task configuration
        """
        self._task[actor_name] = task
        return self

    def add_actor(self, actor_name: str, actor: dict):
        """Add an actor to the scenario

        Args:
            actor_name (str): Name of the actor
            actor (dict): Actor configuration
        """
        self._actor[actor_name] = actor
        return self

    def to_dict(self):
        config = copy.deepcopy(vars(self))
        for key in list(config.keys()):
            if key.startswith("_"):
                val = config.pop(key)
                if hasattr(val, "to_dict"):
                    val = val.to_dict()
                config[key[1:]] = val
        return config
